# language model
import numpy
import theano
import theano.tensor as T
from models import LookupTable, LogisticRegression, GRU, LSTM, Dropout, HiddenLayer, Parser
from algorithm import adadelta
from stream import DStream
import argparse
import logging
import os
import pprint
import numpy as np
import math
import heapq as pq
import cPickle as pkl
import sys
from fuel.datasets import TextFile
from fuel.streams import DataStream
from fuel.transformers import Batch
from fuel.schemes import ConstantScheme
import configurations

logger = logging.getLogger(__name__)

# Get the arguments
parser = argparse.ArgumentParser()
parser.add_argument("--proto",  default="get_config_lm_gru",
                     help="Prototype config to use for config")
args = parser.parse_args()


def topk(arr, K, opset):
    arr = zip(range(0, len(arr[0][0])), arr[0][0])
    arr_reduce = []
    for op in opset:
        arr_reduce.append(arr[op])
    arr = arr_reduce
    arr = sorted(arr, cmp=lambda x,y:cmp(x[1],y[1]), reverse = True)
    return arr[:K]

def decoder(beam_fn, beam_size, word, pos, word_atoi, pos_atoi, nterm_atoi, op_itoa, op_atoi):
    configuration = getattr(configurations, args.proto)()
    sys.path.append("../tools/")
    from binarize import CBinarizedTreeNode as BTN
    count = len(word)
    #print count
    bin_node = BTN()
    qu = []
    for tmp in zip(word, pos):
        qu.append(tmp)
    qu.reverse()
    st = []
    #init
    word, pos, nterm = bin_node.get_feature(qu, st)
    #print word
    #print pos
    #print nterm
    #word = np.matrix([[word_atoi.get(w, 0) for w in word]], dtype = 'int64')
    #pos = np.matrix([[pos_atoi.get(p, 0) for p in pos]], dtype = 'int64')
    #nterm = np.matrix([[nterm_atoi.get(nt, 0) for nt in nterm]], dtype = 'int64')
    #softmax_y = beam_fn(word, pos, nterm, 0)
    #op = topk(softmax_y, beam_size, kid2nt)
    dp = []
    next_qu, next_st = bin_node.one_step('shift', qu, st)
    dp.append([next_qu, next_st, 0, op_atoi['shift']])#cur_qu, cur_st, cur_score, cur_op
    #for action in op:
    #    #print op_itoa[action[0]]
    #    next_qu, next_st = bin_node.one_step(op_itoa[action[0]], qu, st)
    #    if next_qu != None and next_st != None:
    #        dp.append([next_qu, next_st, -math.log(action[1])])
    #dp
    #print dp
    for i in xrange(count * 2 - 2):
        newdp = []
        for [cur_qu, cur_st, cur_score, curop] in dp:
            word, pos, nterm = bin_node.get_feature(cur_qu, cur_st)
            opset = set()
            if len(cur_st) < 2:
                opset.add(op_atoi['shift'])
            else:
                left = cur_st[-2].name + ( '' if cur_st[-2].temporary == False  else  '*' )
                right = cur_st[-1].name + ('' if cur_st[-1].temporary == False else '*' )
                key = left + '#' + right
                #print key
                #if key in kid2nt:
                #if False:
                #if True:
                    #print left + '#' + right
                    #print { op for op in kid2nt[left + '#' + right] }
                    #opset = { op_atoi[op] for op in kid2nt[left + '#' + right] }
                #else:
                opset = set(range(configuration['op_size']))
                    #continue
                if len(cur_qu) >= 1:
                    opset.add(op_atoi['shift'])
            #print word
            #print pos
            #print nterm
            word = np.matrix([[word_atoi.get(w, 0) for w in word]], dtype = 'int64')
            pos = np.matrix([[pos_atoi.get(p, 0) for p in pos]], dtype = 'int64')
            nterm = np.matrix([[nterm_atoi.get(nt, 0) for nt in nterm]], dtype = 'int64')
            softmax_y = beam_fn(word, pos, nterm, 0)
            op = topk(softmax_y, beam_size, opset)
            for action in op:
                #print op_itoa[action[0]]
                next_qu, next_st = bin_node.one_step(op_itoa[action[0]], cur_qu, cur_st)
                if next_qu == None or next_st == None:
                    continue
                #newdp.append([next_qu, next_st, cur_score + math.log(action[1]) + math.log(trans_p[curop][action[0]]), action[0]])
                newdp.append([next_qu, next_st, cur_score + math.log(action[1]), action[0]])
        dp = sorted(newdp, cmp=lambda x,y:cmp(x[2],y[2]), reverse = True)[:beam_size]
    st = ""
    if len(dp) > 0 and len(dp[0]) > 1 and len(dp[0][1]) > 0:
        dp[0][1][0].temporary = False
        #print dp[0][1][0].pretty_print()
        #st =  "(TOP "+ bin_node.unbinary(dp[0][1][0])[0].__str__() + ")"
        st =  bin_node.unbinary(dp[0][1][0])[0].__str__()
        #st =  "(TOP "+ dp[0][1][0].pretty_print() + ")"
    else:
        st = ''
    return st

def parse(path, model):
    configuration = getattr(configurations, args.proto)()
    #logger.info("\nModel options:\n{}".format(pprint.pformat(configuration)))

    use_noise = T.iscalar()

    word = T.lmatrix()

    pos = T.lmatrix()

    nterm = T.lmatrix()

    op = T.lmatrix()

    parser = Parser(**configuration)

    parser.apply(word, pos, nterm, op, use_noise)

    parser.load(model)

    softmax_y = parser.softmax_y
    word_atoi = pkl.load(open(configuration['word_table']))
    pos_atoi = pkl.load(open(configuration['pos_table']))
    nterm_atoi = pkl.load(open(configuration['nterm_table']))
    op_atoi = pkl.load(open(configuration['op_table']))
    op_itoa = pkl.load(open(configuration['op_itoa_table']))

    #init_p = pkl.load(open(configuration['init']))
    #trans_p = pkl.load(open(configuration['trans']))

    #kid2nt = pkl.load(open(configuration['kid2nt']))

    beam_fn = theano.function([word, pos, nterm, use_noise], [softmax_y])

    #decoder

    word_sent_test = open(configuration['word_sent_test'])
    pos_sent_test = open(configuration['pos_sent_test'])
    #for tmp in op_atoi.keys():
    #    print "'" + tmp + "', ", 
    f = open(path, 'w')

    for w, p in zip(word_sent_test, pos_sent_test):
        #print w
        #print p
        w = w.strip().split(' ')
        p = p.strip().split(' ')
        res =  decoder(beam_fn, 2, w, p, word_atoi, pos_atoi, nterm_atoi, op_itoa, op_atoi)
        f.write(res + '\n')

def eval(path, model = ""):
    import commands
    parse(path, model)
    configuration = getattr(configurations, args.proto)()
    res = commands.getoutput(configuration['eval'] + ' ' + configuration['bracket_ref'] + ' ' + path)
    if res == 'nan':
        res = 0
    else:
        res = float(res)
    return res

if __name__=='__main__':
    print eval('./res/aaa')
